
# 参考 https://mujoco.readthedocs.io/en/stable/APIreference/index.html

import time
import math
import numpy as np

import mujoco
import mujoco.viewer

from scipy.spatial.transform import Rotation as R

def quaternion2euler(quaternion):
    r = R.from_quat(quaternion)
    euler = r.as_euler('xyz', degrees=True)
    return euler

m = mujoco.MjModel.from_xml_path('scene.xml')
d = mujoco.MjData(m)


supply = 0.5

# k = [3.8236 0.0099 0.0003]
# x = [摆角，摆速度，电机速度]' 
adjust_ang= 3.8236 / 3;
adjust_p_vel= 0.0099 / 3;
adjust_motor_vel= 0.0003 / 3;
adjust_total_ang= 0;   #微调角度

# LQR稳定控制器功能
# 计算需要设置电机的电压，以稳定摆
def controllerLQR(p_angle, p_vel, m_vel):
    # 如果角度可控
    # 计算控制律 
    # LQR controller u = k*x
    u =  adjust_ang*p_angle + adjust_p_vel*p_vel + adjust_motor_vel*m_vel;

    '''
    # 限制设定给电机的电压
    if(abs(u) > supply):
        u = np.sign(u) * supply;
    '''

    return u;


with mujoco.viewer.launch_passive(m, d) as viewer:
    # Close the viewer automatically after 30 wall-seconds.
    start = time.time()
    cnt = 0

    '''
    def get_sensor_data(sensor_name):
        sensor_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_SENSOR, sensor_name)
        if sensor_id == -1:
            raise ValueError(f"Sensor '{sensor_name}' not found in model!")
        start_idx = m.sensor_adr[sensor_id]
        dim = m.sensor_dim[sensor_id]
        sensor_values = d.sensordata[start_idx : start_idx + dim]
        return sensor_values
    '''

    body_id = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, "bodyskele")

    while viewer.is_running() and time.time() - start < 150:
        step_start = time.time()

        x = d.sensor("orientation").data[1]
        #y = d.sensor("orientation").data[2]
        euler_data = quaternion2euler(d.sensor("orientation").data)
        y_degree = euler_data[1] * math.pi / 180

        '''测试step '''
        d.ctrl[0] = 0.2 * np.sign(x)  #第一个motor设置扭矩

        d.ctrl[1] = controllerLQR(y_degree, d.sensor("body_w").data[1], d.sensor("mom_w").data[1])
        d.ctrl[2] = d.ctrl[1] * 500          # 需按实际需求调整飞轮转速
        '''
        if (abs(y_degree) < 0.5):
            d.ctrl[1] = controllerLQR(y_degree, d.sensor("body_w").data[1], d.sensor("mom_w").data[1])
        else:
            d.ctrl[1] = np.sign(d.sensor("body_w").data[1]) * supply
        '''

        print(f"角度={y_degree}, 线速度={d.sensor("body_w").data[1]}, 飞轮转速={d.sensor("mom_w").data[1]}, 力矩={d.ctrl[1]}")
        mujoco.mj_step(m, d)

        
        #print(f"动量轮倾斜角度{euler_data[1]} 扭矩{d.ctrl[1]}")
        #print(f"机体姿态{d.xquat[body_id]}")
        #print(f"机体传感器姿态{d.sensor("orientation").data}")
        #print(f"飞轮角速度数据{d.sensor("mom_w").data}")


        cnt += 0.01

        # Example modification of a viewer option: toggle contact points every two seconds.
        with viewer.lock():
          viewer.opt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = int(d.time % 2)

        # Pick up changes to the physics state, apply perturbations, update options from GUI.
        viewer.sync()

        # Rudimentary time keeping, will drift relative to wall clock.
        time_until_next_step = m.opt.timestep - (time.time() - step_start)
        if time_until_next_step > 0:
          time.sleep(time_until_next_step)







